import arrow
import traceback

from pymongo import MongoClient, UpdateOne

from conf.settings import CLOUDMONITOR_MONGODB, DATE_FMT
from utility.utils import rollback, backup
from utility.log import logger

# db = MongoClient(CLOUDMONITOR_MONGODB)['cloudmonitor']
db = MongoClient(host='localhost', port=27017)['cloudmonitor']

now_datetime = arrow.utcnow()
reset_day = now_datetime.replace(day=1, hour=2, minute=0, second=0, microsecond=0).shift(months=1, hours=-8).int_timestamp
now_datetime = now_datetime.datetime
now = now_datetime.isoformat()


def main():
    try:
        backup(('balances', 'packages', 'charges'), logger, db, now)
        plan = db['packages'].find_one({'key': 'scanv_custom_a'})
        temp_tuple = ('contentManualTask', 'sslManualTask', 'assetManualTask', 'vulManualTask', 'securityEventManualTask')

        balances_update_list = []
        charges_update_list = []
        for balances in db['balances'].find({'planId': plan['_id']}):
            balances_update_list.append(UpdateOne({'_id': balances['_id']}, {'$inc': {
                'balance.contentManualTask.total': -400,
                'balance.securityEventManualTask.total': -25,
                'balance.sslManualTask.total': -25,
                'balance.assetManualTask.total': -25,
                'balance.vulManualTask.total': -25
            }}))

            for item in temp_tuple:
                query = {'uid': balances['uid'], 'packId': plan['_id'], 'funcType': item, 'chargeType': 'plan'}
                charges_update_list.append(UpdateOne(query, {'$set': {f'total': 100, 'updateTime': now_datetime}}))

        # limit Type
        for charges in db['charges'].find({'funcType': {'$regex': 'ManualTask'}, 'limitType': 'frequency'}):
            charges_update_list.append(UpdateOne({'_id': charges['_id']}, {'$set': {
                'limitType': 'period',
                'time': reset_day,
                'taskType': charges['taskType'].replace('ManualTask', ''),
                'updateTime': now_datetime
            }}))

        # node task type
        for balances in db['balances'].find():
            balances_update_list.append(UpdateOne({'_id': balances['_id']}, {'$set': {
                'balance.httpMaxNode.taskType': 'http',
                'balance.pingMaxNode.taskType': 'ping',
            }}))

        if balances_update_list:
            db['balances'].bulk_write(balances_update_list)
        if charges_update_list:
            db['charges'].bulk_write(charges_update_list)

        db['packages'].update_one({'_id': plan['_id']}, {'$set': {
            'limit.contentManualTask.total': 100,
            'limit.securityEventManualTask.total': 100,
            'limit.sslManualTask.total': 100,
            'limit.assetManualTask.total': 100,
            'limit.vulManualTask.total': 100
        }})

    except Exception as e:
        traceback.print_exc()
        logger.exception(e)
        rollback(('balances', 'packages', 'charges'), logger, db, now)


if __name__ == '__main__':
    main()

